import pandas as pd
import numpy as np
import os
import matplotlib.pyplot as plt
import matplotlib as mpl
import statsmodels.api as sm
from statsmodels.tsa.statespace.sarimax import SARIMAX
from statsmodels.tsa.arima.model import ARIMA
import warnings
from scipy import stats
import calendar
import matplotlib.dates as mdates

# 设置中文显示
mpl.rcParams['font.sans-serif'] = ['SimHei']  # 使用黑体
mpl.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

warnings.filterwarnings('ignore')

EXCEL_PATH = "E:/M/MathModel/23国赛C题/ds/历年6月底到7月初.xlsx"
OUTPUT_DIR = "E:/M/MathModel/23国赛C题/ds/销售分析图表"


def load_and_prepare_data(file_path):
    # 读取Excel中的四个表
    df_2020 = pd.read_excel(file_path, sheet_name='2020')
    df_2021 = pd.read_excel(file_path, sheet_name='2021')
    df_2022 = pd.read_excel(file_path, sheet_name='2022')
    df_2023 = pd.read_excel(file_path, sheet_name='2023')

    # 合并数据
    df = pd.concat([df_2020, df_2021, df_2022, df_2023])

    # 转换日期格式
    df['销售日期'] = pd.to_datetime(df['销售日期'])
    df['month_day'] = df['销售日期'].dt.strftime('%m-%d')


    return df


def create_daily_dataset(df):
    # 按'月-日'和品类分类
    daily_data = df.groupby(['month_day', '分类名称']).agg(
        total_sales=('销量(千克)', 'mean'),
        avg_loss_rate=('损耗率(%)', 'mean'),
        data_points=('销量(千克)', 'count')
    ).reset_index()

    # 创建表
    sales_pivot = daily_data.pivot(
        index='month_day',
        columns='分类名称',
        values='total_sales'
    )

    loss_pivot = daily_data.pivot(
        index='month_day',
        columns='分类名称',
        values='avg_loss_rate'
    )

    # 添加日期特征
    sales_pivot['date'] = pd.to_datetime('2023-' + sales_pivot.index)
    sales_pivot.set_index('date', inplace=True)

    # 添加时间特征
    sales_pivot['day_of_week'] = sales_pivot.index.dayofweek
    sales_pivot['is_weekend'] = sales_pivot['day_of_week'].apply(lambda x: 1 if x >= 5 else 0)
    sales_pivot['day_of_year'] = sales_pivot.index.dayofyear

    return sales_pivot, loss_pivot


def forecast_category(category, data, category_type):
    # 检查数据完整性
    sales_col = f'sales_{category}'
    if sales_col not in data.columns:
        print(f"警告: 品类 '{category}' 无销售数据，跳过预测")
        return None, None, None

    # 确保有足够的数据点
    sales_data = data[sales_col].dropna()
    print(f"{category}数据点数量: {len(sales_data)}")

    if len(sales_data) < 10:
        print(f"警告: {category} 只有 {len(sales_data)} 个数据点，可能不足以建模")

    # 根据品类类型选择模型
    if category_type == '花叶类':
        # ARIMA模型 (1,1,1)
        try:
            model = ARIMA(
                sales_data,
                order=(1, 1, 1)
            )
            model_fit = model.fit()
            forecast = model_fit.forecast(steps=7)

            # 确保预测结果是Series类型
            if not isinstance(forecast, pd.Series):
                forecast = pd.Series(forecast, index=pd.date_range('2023-07-01', periods=7, freq='D'))

            return model_fit, forecast, None
        except Exception as e:
            print(f"ARIMA模型拟合失败: {e}")
            # 使用简单预测
            mean_value = sales_data.mean()
            forecast = pd.Series([mean_value] * 7, index=pd.date_range('2023-07-01', periods=7, freq='D'))
            return None, forecast, None
    else:
        # SARIMA模型
        try:
            model = SARIMAX(
                sales_data,
                order=(0, 1, 1),
                seasonal_order=(0, 1, 1, 7),
                exog=data[['is_weekend']]
            )
            model_fit = model.fit(disp=False)

            # 创建未来7天的外生变量
            future_dates = pd.date_range('2023-07-01', periods=7, freq='D')
            future_exog = pd.DataFrame({
                'is_weekend': [1 if d.weekday() >= 5 else 0 for d in future_dates],
                'date': future_dates
            }).set_index('date')

            forecast_obj = model_fit.get_forecast(steps=7, exog=future_exog)
            forecast = forecast_obj.predicted_mean
            conf_int = forecast_obj.conf_int()

            # 确保预测结果是Series类型
            if not isinstance(forecast, pd.Series):
                forecast = pd.Series(forecast, index=future_dates)

            return model_fit, forecast, conf_int
        except Exception as e:
            print(f"SARIMA模型拟合失败: {e}")
            # 使用简单预测
            mean_value = sales_data.mean()
            forecast = pd.Series([mean_value] * 7, index=pd.date_range('2023-07-01', periods=7, freq='D'))
            return None, forecast, None


def generate_model_report(model_fit, category, model_type, output_dir):
    # 创建报告目录
    report_dir = os.path.join(output_dir, "model_reports")
    os.makedirs(report_dir, exist_ok=True)

    # 文本报告
    report_path = os.path.join(report_dir, f"{category}_{model_type}_report.txt")
    with open(report_path, 'w') as f:
        f.write(f"{'=' * 50}\n")
        f.write(f"{category} {model_type} 模型参数报告\n")
        f.write(f"{'=' * 50}\n\n")

        # 添加模型摘要
        try:
            f.write(model_fit.summary().as_text())
        except Exception as e:
            f.write(f"无法生成模型摘要: {e}\n")

        # 添加模型公式
        if model_type == 'ARIMA':
            f.write("\n\n模型公式:\n")
            try:
                order = model_fit.model.order
                f.write(f"ARIMA{order}\n")
            except:
                f.write("无法获取模型公式\n")
        else:
            f.write("\n\n模型公式:\n")
            try:
                order = model_fit.model.order
                seasonal_order = model_fit.model.seasonal_order
                f.write(f"SARIMA{order}x{seasonal_order}\n")
            except:
                f.write("无法获取模型公式\n")

    # 参数可视化
    plt.figure(figsize=(12, 6))

    try:
        if model_type == 'ARIMA':
            # ARIMA参数图
            params = model_fit.params
            std_err = model_fit.bse
            param_names = params.index

            plt.bar(param_names, params, yerr=std_err, capsize=5)
            plt.title(f"{category} ARIMA 模型参数估计")
            plt.ylabel("参数值")
            plt.grid(axis='y', linestyle='--', alpha=0.7)
        else:
            # SARIMA参数图
            param_table = pd.DataFrame({
                '参数': model_fit.param_names,
                '估计值': model_fit.params,
                '标准误': model_fit.bse,
                'z值': model_fit.tvalues,
                'P>|z|': model_fit.pvalues
            })

            # 显著参数突出显示
            significant = param_table['P>|z|'] < 0.05
            colors = ['green' if sig else 'gray' for sig in significant]

            plt.bar(param_table['参数'], param_table['估计值'],
                    yerr=param_table['标准误'], capsize=5,
                    color=colors)
            plt.title(f"{category} SARIMA 模型参数估计")
            plt.ylabel("参数值")
            plt.xticks(rotation=45, ha='right')
            plt.grid(axis='y', linestyle='--', alpha=0.7)

        plt.tight_layout()
        plt.savefig(os.path.join(report_dir, f"{category}_parameters.png"))
        plt.close()
    except Exception as e:
        print(f"参数可视化失败: {e}")

    # 残差诊断图
    try:
        plt.figure(figsize=(12, 8))
        model_fit.plot_diagnostics(figsize=(12, 8))
        plt.suptitle(f"{category} 模型诊断图", fontsize=16)
        plt.tight_layout()
        plt.savefig(os.path.join(report_dir, f"{category}_diagnostics.png"))
        plt.close()
    except Exception as e:
        print(f"无法为 {category} 生成诊断图: {e}")

    return report_path


def calculate_replenishment(forecast, loss_rate):
    if loss_rate >= 100:  # 防止损耗率异常
        loss_rate = 30
    return forecast / (1 - loss_rate / 100)


def main():
    # 创建输出目录
    os.makedirs(OUTPUT_DIR, exist_ok=True)

    # 加载数据
    print("加载数据")
    try:
        df = load_and_prepare_data(EXCEL_PATH)
        print(f"加载成功，共 {len(df)} 条记录")
    except Exception as e:
        print(f"数据加载失败: {e}")
        return

    # 创建数据集
    print("创建每日数据集")
    try:
        sales_data, loss_data = create_daily_dataset(df)
        print(f"创建成功，共 {len(sales_data)} 天数据")

        # 重命名列以明确区分
        sales_data.columns = [f'sales_{col}' if col not in ['day_of_week', 'is_weekend', 'day_of_year'] else col
                              for col in sales_data.columns]
    except Exception as e:
        print(f"数据集创建失败: {e}")
        return

    # 定义要预测的品类
    categories = {
        '花叶类': 'ARIMA',
        '花菜类': 'SARIMA',
        '辣椒类': 'SARIMA',
        '茄类': 'SARIMA',
        '水生根茎类': 'SARIMA',
        '食用菌': 'SARIMA'
    }

    # 预测每个品类
    print("开始预测各品类销量")
    forecasts = {}
    conf_intervals = {}
    replenishments = {}
    model_reports = {}
    avg_loss_rates = {}

    for category, model_type in categories.items():
        # 检查品类是否存在
        if f'sales_{category}' not in sales_data.columns:
            print(f"跳过 '{category}' - 无数据")
            continue

        print(f"  - 预测 {category} ({model_type})...")

        # 获取平均损耗率
        if category in loss_data.columns:
            avg_loss_rate = loss_data[category].mean()
        else:
            print(f"警告: {category} 无损耗率数据，使用默认值10%")
            avg_loss_rate = 10.0
        avg_loss_rates[category] = avg_loss_rate

        # 预测销量
        model_fit, sales_forecast, conf_int = forecast_category(category, sales_data, model_type)

        if sales_forecast is None:
            print(f"  !! {category} 预测失败")
            continue

        forecasts[category] = sales_forecast
        conf_intervals[category] = conf_int

        # 计算补货量（考虑损耗率）
        replenishment = calculate_replenishment(sales_forecast, avg_loss_rate)
        replenishments[category] = replenishment

        # 生成模型报告
        if model_fit is not None:
            try:
                report_path = generate_model_report(model_fit, category, model_type, OUTPUT_DIR)
                model_reports[category] = report_path
                print(f"    模型报告保存至: {report_path}")
            except Exception as e:
                print(f"    模型报告生成失败: {e}")
        else:
            print("    无模型对象，跳过报告生成")

    # 创建预测结果DataFrame
    forecast_dates = pd.date_range('2023-07-01', periods=7, freq='D')

    # 确保所有预测都是Series类型
    for cat in forecasts:
        if not isinstance(forecasts[cat], pd.Series):
            forecasts[cat] = pd.Series(forecasts[cat], index=forecast_dates)

    # 创建DataFrame
    forecast_df = pd.DataFrame(forecasts)

    # 同样处理补货量
    for cat in replenishments:
        if not isinstance(replenishments[cat], pd.Series):
            replenishments[cat] = pd.Series(replenishments[cat], index=forecast_dates)

    replenishment_df = pd.DataFrame(replenishments)

    # 保存结果
    print("保存结果")
    try:
        # 保存为Excel
        with pd.ExcelWriter(os.path.join(OUTPUT_DIR, 'forecast_results.xlsx')) as writer:
            if not forecast_df.empty:
                forecast_df.to_excel(writer, sheet_name='销量预测')
            else:
                pd.DataFrame().to_excel(writer, sheet_name='销量预测')

            if not replenishment_df.empty:
                replenishment_df.to_excel(writer, sheet_name='补货量建议')
            else:
                pd.DataFrame().to_excel(writer, sheet_name='补货量建议')

            # 添加损耗率信息
            if avg_loss_rates:
                loss_rate_df = pd.DataFrame({
                    '品类': list(avg_loss_rates.keys()),
                    '平均损耗率(%)': list(avg_loss_rates.values())
                })
                loss_rate_df.to_excel(writer, sheet_name='损耗率信息', index=False)
            else:
                pd.DataFrame(columns=['品类', '平均损耗率(%)']).to_excel(writer, sheet_name='损耗率信息', index=False)

        print("Excel文件保存成功")
    except Exception as e:
        print(f"Excel保存失败: {e}")

    # 可视化结果
    print("生成可视化图表")
    try:
        if forecasts:
            plt.figure(figsize=(14, 10))
            for i, category in enumerate(forecasts.keys(), 1):
                plt.subplot(3, 2, i)

                # 获取历史数据
                hist_col = f'sales_{category}'
                if hist_col in sales_data.columns:
                    hist_data = sales_data[hist_col].dropna()
                    hist_data.plot(label='历史销量', marker='o', alpha=0.7)

                # 绘制预测数据
                forecasts[category].plot(label='预测销量', marker='s', linestyle='--', color='red')

                # 添加置信区间（SARIMA）
                if category in conf_intervals and conf_intervals[category] is not None:
                    conf_int = conf_intervals[category]
                    plt.fill_between(
                        conf_int.index,
                        conf_int.iloc[:, 0],
                        conf_int.iloc[:, 1],
                        color='red',
                        alpha=0.1,
                        label='95%置信区间'
                    )

                # 格式设置
                plt.title(f'{category}销量预测')
                plt.xlabel('日期')
                plt.ylabel('销量(千克)')
                plt.gca().xaxis.set_major_formatter(mdates.DateFormatter('%m-%d'))
                plt.legend()
                plt.grid(True, linestyle='--', alpha=0.7)

            plt.tight_layout()
            plt.savefig(os.path.join(OUTPUT_DIR, 'sales_forecast_comparison.png'))
            print("销量预测图保存成功")
        else:
            print("无预测数据，跳过销量预测图生成")
    except Exception as e:
        print(f"销量预测图生成失败: {e}")

    # 补货建议可视化
    try:
        if replenishments:
            plt.figure(figsize=(14, 8))
            ax = plt.subplot(111)
            width = 0.12  # 条形宽度
            dates = replenishment_df.index

            for i, (category, values) in enumerate(replenishments.items()):
                positions = np.arange(len(dates)) + i * width
                plt.bar(positions, values, width=width, label=category)

            # 格式设置
            plt.title('2023年7月1-7日补货建议')
            plt.xlabel('日期')
            plt.ylabel('补货量(千克)')
            plt.xticks(np.arange(len(dates)) + width * 2.5, [d.strftime('%m-%d') for d in dates])
            plt.legend()
            plt.grid(True, axis='y', linestyle='--', alpha=0.7)
            plt.tight_layout()
            plt.savefig(os.path.join(OUTPUT_DIR, 'replenishment_suggestions.png'))
            print("补货建议图保存成功")
        else:
            print("无补货建议数据，跳过补货建议图生成")
    except Exception as e:
        print(f"补货建议图生成失败: {e}")

    # 输出最终结果
    print("\n2023年7月1-7日预测结果")
    if forecasts:
        print("销量预测:")
        print(forecast_df.round(2))
    else:
        print("无销量预测数据")

    if replenishments:
        print("\n补货建议(考虑损耗率):")
        print(replenishment_df.round(2))
    else:
        print("\n无补货建议数据")

    if avg_loss_rates:
        print("\n平均损耗率:")
        for category, rate in avg_loss_rates.items():
            print(f"  - {category}: {rate:.2f}%")
    else:
        print("\n无损耗率数据")

    print(f"\n所有结果已保存至: {os.path.abspath(OUTPUT_DIR)}")


if __name__ == "__main__":
    main()